{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "\n",
    "import os\n",
    "# from utils.load_smplx_params import load_multiperson_smplx_params\n",
    "import numpy as np\n",
    "import pickle\n",
    "import math\n",
    "import cv2\n",
    "from smplx import smplx\n",
    "import torch\n",
    "from nosmpl.vis.vis_o3d import vis_mesh_o3d\n",
    "import open3d as o3d\n",
    "from utils.process_transformation import np_mat2axangle, torch_mat2axangle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 23, 3, 3])\n",
      "tensor([[-1.1035,  2.0479,  1.2528,  0.4798,  0.4667,  0.1900,  0.0172, -0.4207,\n",
      "          0.0557,  0.3166]], device='cuda:0', requires_grad=True)\n",
      "torch.Size([1, 1, 3, 3])\n"
     ]
    }
   ],
   "source": [
    "f = open(\"/home/liuyun/HHO-dataset/data_processing/smplx/smplx_output/20230805_1/beta/p1/T_Pose_x.pkl\", 'rb')\n",
    "smpl = pickle.load(f)\n",
    "\n",
    "print(smpl['body_pose'].shape)\n",
    "print(smpl['betas'])\n",
    "print(smpl['global_orient'].shape)\n",
    "\n",
    "betas = smpl['betas'].detach().cpu()\n",
    "\n",
    "smpl_model = smplx.create(\"/home/liuyun/HHO-dataset/data_processing/smplx/transfer_data\", model_type='smpl')\n",
    "\n",
    "# print(torch_mat2axangle(smpl['body_pose'].detach().cpu().squeeze()).shape)\n",
    "# print(torch_mat2axangle(smpl['global_orient'].detach().cpu().squeeze().unsqueeze(0)).shape)\n",
    "\n",
    "output = smpl_model(betas = smpl['betas'].detach().cpu(), \n",
    "                        body_pose = torch_mat2axangle(smpl['body_pose'].detach().cpu().squeeze()).unsqueeze(0),\n",
    "                        global_orient = torch_mat2axangle(smpl['global_orient'].detach().cpu().squeeze().unsqueeze(0)).unsqueeze(0),\n",
    "                        return_verts=True)\n",
    "\n",
    "vertices = output.vertices[0].detach().cpu().numpy().squeeze()\n",
    "joints = output.joints[0].detach().cpu().numpy().squeeze()\n",
    "\n",
    "faces = smpl_model.faces.astype(np.int32)\n",
    "mesh = o3d.geometry.TriangleMesh()\n",
    "mesh.vertices = o3d.utility.Vector3dVector(vertices)\n",
    "mesh.triangles = o3d.utility.Vector3iVector(faces)\n",
    "o3d.io.write_triangle_mesh(\"T_Pose_smpl.ply\", mesh)\n",
    "\n",
    "f.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "T Pose for smpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.load(\"/share/datasets/HHO_dataset/data/20230805_1/000/SMPLX_fitting/person_1/150to199.npz\", allow_pickle=True)\n",
    "data = x['results'].item()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "smpl_model = smplx.create(\"/home/liuyun/HHO-dataset/data_processing/smplx/transfer_data\", model_type='smpl')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "smpl_body_pose = np.full((50, 23, 3), 0, dtype=np.float32) \n",
    "smpl_body_pose = torch.from_numpy(smpl_body_pose)\n",
    "global_orient = np.full((3), 0, dtype=np.float32)\n",
    "global_orient = torch.from_numpy(global_orient)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = smpl_model(betas = data['betas'][0].unsqueeze(0).cpu(), \n",
    "                        body_pose = smpl_body_pose[0].unsqueeze(0),\n",
    "                        global_orient = global_orient.unsqueeze(0).unsqueeze(0),\n",
    "                        return_verts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertices = output.vertices[0].detach().cpu().numpy().squeeze()\n",
    "joints = output.joints[0].detach().cpu().numpy().squeeze()\n",
    "\n",
    "faces = smpl_model.faces.astype(np.int32)\n",
    "mesh = o3d.geometry.TriangleMesh()\n",
    "mesh.vertices = o3d.utility.Vector3dVector(vertices)\n",
    "mesh.triangles = o3d.utility.Vector3iVector(faces)\n",
    "o3d.io.write_triangle_mesh(\"T_Pose.ply\", mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "T Pose for smplx\n",
    "wrong for \n",
    "\n",
    "print(data['betas'][0].unsqueeze(0).cpu().shape)\n",
    "print(torch.randn([1, smplx_model.num_betas], dtype=torch.float32).shape)\n",
    "\n",
    "torch.Size([1, 10])\n",
    "torch.Size([1, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 10])\n",
      "torch.Size([1, 10])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.load(\"/share/datasets/HHO_dataset/data/20230805_1/000/SMPLX_fitting/person_1/150to199.npz\", allow_pickle=True)\n",
    "data = x['results'].item()\n",
    "\n",
    "smplx_model = smplx.create(\"/share/human_model/models\", model_type='smplx', num_betas=10)\n",
    "\n",
    "# smpl_body_pose = np.full((50, 21, 3), 0, dtype=np.float32) \n",
    "# smpl_body_pose = torch.from_numpy(smpl_body_pose)\n",
    "\n",
    "expression = torch.zeros([1, smplx_model.num_expression_coeffs], dtype=torch.float32)\n",
    "body_pose = torch.zeros([1, 21, 3], dtype=torch.float32)\n",
    "\n",
    "print(data['betas'][0].unsqueeze(0).cpu().shape)\n",
    "print(torch.randn([1, smplx_model.num_betas], dtype=torch.float32).shape)\n",
    "output = smplx_model(\n",
    "        betas = data['betas'][0].unsqueeze(0).cpu(), expression=expression, body_pose=body_pose, return_verts=True\n",
    "    )\n",
    "\n",
    "vertices = output.vertices[0].detach().cpu().numpy().squeeze()\n",
    "joints = output.joints[0].detach().cpu().numpy().squeeze()\n",
    "\n",
    "faces = smplx_model.faces.astype(np.int32)\n",
    "mesh = o3d.geometry.TriangleMesh()\n",
    "mesh.vertices = o3d.utility.Vector3dVector(vertices)\n",
    "mesh.triangles = o3d.utility.Vector3iVector(faces)\n",
    "o3d.io.write_triangle_mesh(\"T_Pose_x.ply\", mesh)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fit beta from smplx to smpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from .smplx import build_layer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_cfg = \"./smplx/config_files/smplx2smpl.yaml\"\n",
    "data_dir = \"/share/datasets/HHO_dataset/data/20230806_1\"\n",
    "device = torch.device(\"cuda:0\")\n",
    "model_path = exp_cfg.body_model.folder\n",
    "body_model = build_layer(model_path, **exp_cfg.body_model)\n",
    "\n",
    "body_model = body_model.to(device=device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "84\n",
      "628\n",
      "['/share/hhodataset/VTS/20231002/062', '/share/hhodataset/VTS/20231003_1/030', '/share/hhodataset/VTS/20231003_1/021', '/share/hhodataset/VTS/20231003_2/034', '/share/hhodataset/VTS/20231003_2/020', '/share/hhodataset/VTS/20231008/053', '/share/hhodataset/VTS/20231011/050', '/share/hhodataset/VTS/20231011/089', '/share/hhodataset/VTS/20231020/072', '/share/hhodataset/VTS/20231030/064']\n"
     ]
    }
   ],
   "source": [
    "# 统计 有多少deleted\n",
    "sets = [\"20231001\", \"20231002\", \"20231003_1\", \"20231003_2\", \"20231008\", \"20231011\", \"20231018\", \"20231020\", \"20231023\", \"20231030\", \"20231108\"]\n",
    "dataset_dir = \"/share/hhodataset/VTS/\"\n",
    "# /share/hhodataset/VTS/20231020/001/SMPLX_fitting\n",
    "deleted = 0\n",
    "available = 0\n",
    "big = []\n",
    "for s in sets:\n",
    "    set_dir = os.path.join(dataset_dir, s)\n",
    "    for d in os.listdir(set_dir):\n",
    "        if os.path.exists(os.path.join(set_dir, d, \"SMPLX_fitting\")):\n",
    "            if not os.path.exists(os.path.join(set_dir, d, \"deleted\")):\n",
    "                available += 1\n",
    "                if os.path.exists(os.path.join(set_dir, d, \"big\")):\n",
    "                    big.append(os.path.join(set_dir, d))                \n",
    "            else:\n",
    "                deleted += 1\n",
    "                \n",
    "print(deleted)\n",
    "print(available)\n",
    "print(big)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "HHO",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
